from langchain.document_loaders import TextLoader, PyPDFLoader, Docx2txtLoader
import os

def load_documents_from_folder(folder_path):
    """
    加载指定文件夹下所有支持的文档（txt, pdf, docx），返回文档对象列表
    """
    docs = []
    for filename in os.listdir(folder_path):
        file_path = os.path.join(folder_path, filename)
        if filename.endswith('.txt'):
            loader = TextLoader(file_path, encoding='utf-8')
            docs.extend(loader.load())
        elif filename.endswith('.pdf'):
            loader = PyPDFLoader(file_path)
            docs.extend(loader.load())
        elif filename.endswith('.docx'):
            loader = Docx2txtLoader(file_path)
            docs.extend(loader.load())
        # 你可以根据需要添加更多格式
    return docs

if __name__ == "__main__":
    folder = "./your_docs_folder"  # 替换为你的文档文件夹路径
    documents = load_documents_from_folder(folder)
    print(f"共加载{len(documents)}个文档片段")
    for doc in documents[:3]:  # 只打印前3个示例
        print(doc)
